# -*- coding: utf-8 -*-
'''
Created on 14.12.2019

@author: yu03
'''
import glob
import re
import datetime
import cv2
import matplotlib.pyplot as plt
from scipy import signal
from FFT_Interpolation import FFT_interpolation_2, line_cal, line_cal_fix
import numpy as np
import os
from scipy.optimize import curve_fit
from mpl_toolkits.mplot3d import Axes3D
from PyUeye_Unified.Cross_200line_SAAC import folder_path, np_names, hor_index, hor_lines, ver_lines, line_num
import sys
from tqdm import tqdm

# calculation_mode = 'diff' ### DC removed
calculation_mode = 'DC' ### DC Not removed


''' 
    General Parameters
'''
fs_cam = 25
Lamda = 633e-9
pix_size = 5.3e-6
V_x, V_y, V_z = 0, 0, 0
print(np_names)
# np_names = [r'C:\Users\yu03\Desktop\Fine_step_test\measuement_1217-1406.npy']
for np_file in np_names:
    ''' 
        File name define
    '''
    file_name = re.findall(r"Drift\\(.+).npy", np_file)[0]
    if calculation_mode == 'DC':
        np_result_name = folder_path + '\\results_100line_Gaus\\' + file_name + '_lines' + '.npy'
    elif calculation_mode == 'diff':
        np_result_name = folder_path + '\\results_100line_Gaus\\diff\\' + file_name + '_lines_diff' + '.npy'
    else:
        sys.exit('FFT Mode Error:\n calculation_mode = %s'%calculation_mode)
    print('File:', file_name, '; ', datetime.datetime.now().time())
#     print(np_result_name)
    ''' 
        Reading .npy file
    '''
    f = open(np_file, 'rb')
    f_size = os.fstat(f.fileno()).st_size
    frames = []
    hor_lines_all = []
    ver_lines_all = []
    time_sequence = []
    while f.tell() < f_size:
    #         print(f.tell())
        frame = np.load(f, allow_pickle=True)
        hor_lines_all.append(frame[0])
        ver_lines_all.append(frame[1])
        time_sequence.append(frame[2])    
    # x_axis = np.linspace(0, len(time_sequence)-1, len(time_sequence))
    time_sequence = np.array(time_sequence)
    frames = np.array(frames)
    hor_lines_all = np.array(hor_lines_all)
    ver_lines_all = np.array(ver_lines_all)
    hor_lines_1frame = hor_lines_all[0]
    ver_lines_1frame = ver_lines_all[0]
#     print(np.shape(hor_lines_all))
    frame_total = len(hor_lines_all)
#     line_num = len(hor_lines_1frame)
    pixel_num = len(hor_lines_1frame[0])
    print(frame_total, line_num, '+', line_num, pixel_num)
#     print(hor_lines_all[::,::,580:600])
#     print(ver_lines_all[::,::,580:600])
    if 0:
        plt.figure(1)
        plt.subplot(1,2,1)
        for j in hor_lines_all[::25,0]:
            plt.plot(j)
        plt.subplot(1,2,2)
        for j in ver_lines_all[50,::25]:
            plt.plot(j)
        plt.show()
#         plt.savefig(folder_path + '\\' + file_name + '.png', dpi=300)
#         plt.close()
    if 1: ###
        ''' 
            Finding Phase Spectrum Component
        '''
        hor_m_k_num_set = []
        ver_m_k_num_set = []
        for j in range(frame_total)[::5]:
            hor_center = np.array(hor_lines_all[j,0])
            ver_center = np.array(ver_lines_all[j,0])
            #hor_center = np.array([i[0]+256*i[1] for i in hor_center])
            if calculation_mode == 'diff':
                hor_center = np.diff(hor_center.astype('int'))
                hor_center = np.concatenate(([0], hor_center))
            elif calculation_mode == 'DC':
                pass
            else:
                sys.exit('FFT Mode Error:\n calculation_mode = %s'%calculation_mode)
              
            hor_f_fit, hor_phase_estim, hor_m_k_num = line_cal(hor_center)
            ver_f_fit, ver_phase_estim, ver_m_k_num = line_cal(ver_center)
            hor_m_k_num_set.append(hor_m_k_num)
            ver_m_k_num_set.append(ver_m_k_num)
      
        m_k_num_hor = int(np.average(hor_m_k_num_set))
        m_k_num_ver = int(np.average(ver_m_k_num_set))
        print(frame_total, "Frames;  ", 'Chosen Components:', m_k_num_hor, m_k_num_ver, '  ', datetime.datetime.now().time() )
    #     
        ''' 
            Processing
        '''    
        hor_f_set = []
        hor_phi_set = []
        ver_f_set = []
        ver_phi_set = []
        frame_num = 0
    #     time_sequence = time_sequence[::100]
    #     for j in range((frame_total+99)//100):
        for j in (range(frame_total)):
            frame_num += 1
            for i in range(line_num):
                hor_center = np.array(hor_lines_all[j,i])
                ver_center = np.array(ver_lines_all[j,i])
    
                if calculation_mode == 'diff':
                    hor_center = np.diff(hor_center.astype('int'))
                    hor_center = np.concatenate(([0], hor_center))
                    ver_center = np.diff(ver_center.astype('int'))
                    ver_center = np.concatenate(([0], ver_center))
                elif calculation_mode == 'DC':
                    pass
                else:
                    sys.exit('FFT Mode Error:\n calculation_mode = %s'%calculation_mode)
                 
                hor_f_fit, hor_phase_estim, hor_m_k_num = line_cal_fix(hor_center, m_k_num_hor, m_k_num_ver)
                ver_f_fit, ver_phase_estim, ver_m_k_num = line_cal_fix(ver_center, m_k_num_hor, m_k_num_ver)
                 
                hor_f_set.append(hor_f_fit)
                hor_phi_set.append(hor_phase_estim)
                ver_f_set.append(ver_f_fit)
                ver_phi_set.append(ver_phase_estim)
            if frame_num%100 == 0:
                print('Frame: ', frame_num, '   ',datetime.datetime.now().time())   
        hor_f_set = np.array(hor_f_set)
        ver_f_set = np.array(ver_f_set)
        hor_f_set = hor_f_set.reshape(line_num, frame_num, order='F')
        ver_f_set = ver_f_set.reshape(line_num, frame_num, order='F')
      
        hor_phi_set = np.array(hor_phi_set)
        ver_phi_set = np.array(ver_phi_set)
        hor_phi_set = hor_phi_set.reshape(line_num, frame_num, order='F')
        ver_phi_set = ver_phi_set.reshape(line_num, frame_num, order='F')
      
        hor_angle_set = (V_x-Lamda*hor_f_set/2)*1e6 ### urad
        ver_angle_set = (V_x-Lamda*ver_f_set/2)*1e6 ### urad
        hor_length_set = np.unwrap(hor_phi_set)/4/np.pi*Lamda*1e9 ### nm
        ver_length_set = np.unwrap(ver_phi_set)/4/np.pi*Lamda*1e9 ### nm
         
        print('line, frame:',np.shape(hor_angle_set))
        f = open(np_result_name,'ab')
        np.save(f, hor_angle_set, allow_pickle=True)
        np.save(f, ver_angle_set, allow_pickle=True)
        np.save(f, hor_length_set, allow_pickle=True)
        np.save(f, ver_length_set, allow_pickle=True)
        np.save(f, time_sequence, allow_pickle=True)
        print('Saved')
        print('')
      
        ''' 
            Plotting
        '''   
        y = [np.arange(frame_num)]*line_num
        y = np.array(y)
        
        fig = plt.figure('Raw')
        ax1 = fig.add_subplot(2, 2, 1, projection='3d') 
        ax2 = fig.add_subplot(2, 2, 2, projection='3d') 
        ax3 = fig.add_subplot(2, 2, 3, projection='3d') 
        ax4 = fig.add_subplot(2, 2, 4, projection='3d') 
        plt.gca().patch.set_facecolor('white')
        x = [hor_lines]*frame_num
        x = np.array(x).T
        for i in range(line_num):
            ax1.plot(x[i,:],y[i,:],hor_angle_set[i,:])
        x = [ver_lines]*frame_num
        x = np.array(x).T
        for i in range(line_num):
            ax2.plot(x[i,:],y[i,:],ver_angle_set[i,:])
        x = [hor_lines]*frame_num
        x = np.array(x).T
        for i in range(line_num):
            ax3.plot(x[i,:],y[i,:],hor_length_set[i,:])
        x = [hor_lines]*frame_num
        x = np.array(x).T
        for i in range(line_num):
            ax4.plot(x[i,:],y[i,:],ver_length_set[i,:])
            
        ax1.set_title('hor_tilt')
        ax2.set_title('ver_tilt')
        ax3.set_title('hor_length')
        ax4.set_title('ver_length')
        
        plt.tight_layout()
        plt.savefig(folder_path + '\\' + file_name + '_raw.png', dpi=300)
        plt.close()
    
#     plt.show()